import  os
import  numpy as np
from sklearn.neighbors import KNeighborsClassifier

def readFileName(dirname):
    x = []
    y = []
    filenames = list_filename(dirname)
    for filename in filenames:
        #循环遍历所有文件名
        arr = filename.split("/")
        y.append(eval(arr[1][0]))
        with open(filename,"rb") as fp:
            str = ""
            #逐一打开文件
            while True:
                line  = fp.readline().decode("utf-8").strip();
                if line:
                    str += line
                else:
                    x.append(list(str))
                    break
    return np.array(x,dtype="int"),y

def list_filename(dirname):
    listFileName = os.listdir(dirname)
    filenames = [os.path.join(dirname+"/",filename) for filename in listFileName]
    return filenames

if __name__ == '__main__':
    x_train,y_train = readFileName("TrainData")
    x_test,y_test = readFileName("TestData")
    estimate = KNeighborsClassifier(n_neighbors=5)
    estimate.fit(x_train,y_train)
    print("模型得分：\n",estimate.score(x_test,y_test))
    y_pred = estimate.predict(x_test)
    print("预测值与真实值对比：\n",np.array(y_test)==y_pred)